from flask import Flask, request, jsonify
from transformers import pipeline

app = Flask(__name__)

# 使用 zero-shot 分类器，指定我们关心的两个类别
classifier = pipeline(
    "zero-shot-classification",
    model="distilbert-base-uncased"
)

@app.route("/analyze-log", methods=["POST"])
def analyze():
    # 获取 JSON 请求体中的日志内容
    data = request.get_json()
    log = data.get("log", "").strip()
    
    if not log:
        return jsonify({"error": "Missing 'log' field in request"}), 400

    # 定义候选标签
    candidate_labels = ["normal log", "anomalous log"]

    # 模型推理
    result = classifier(log, candidate_labels)

    # 提取最高分的类别和置信度
    top_label = result["labels"][0]
    confidence = result["scores"][0]

    # 判断是否为异常
    is_anomalous = top_label == "anomalous log"

    # 可扩展：简单根因推测（可替换为更复杂逻辑或知识库）
    # root_cause = "Unknown"
    # if "timeout" in log.lower():
    #     root_cause = "Network timeout detected"
    # elif "memory" in log.lower() or "oom" in log.lower():
    #     root_cause = "Out of memory"
    # elif "failed" in log.lower() or "error" in log.lower():
    #     root_cause = "Operation failed"
    # 更多规则可继续添加...

    return jsonify({
        "classification": "ANOMALOUS" if is_anomalous else "NORMAL",
        "confidence": float(confidence),
        "is_anomalous": is_anomalous,
        "top_match": top_label,
        "all_scores": dict(zip(result["labels"], result["scores"])),  # 可选：返回全部得分
        "root_cause": root_cause
    })

# 健康检查接口
@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "OK", "model_ready": True})

if __name__ == "__main__":
    print("🚀 日志分析服务启动中... 请访问 http://localhost:5000/health 测试")
    app.run(host="127.0.0.1", port=5000, debug=False)  # 生产环境建议关闭 debug
